package org.luxor.commons.selenium.client.impl;

import cn.hutool.core.thread.ThreadUtil;
import org.luxor.commons.selenium.client.SeleniumClient;
import org.luxor.commons.selenium.component.Context;
import org.luxor.commons.selenium.component.ExecState;
import org.luxor.commons.selenium.component.SeleniumDriver;
import org.luxor.commons.selenium.component.SeleniumDriverCreator;
import org.luxor.commons.selenium.component.command.SeleniumCommand;
import org.luxor.commons.selenium.config.properties.SeleniumProperties;
import org.luxor.commons.selenium.exception.CreateWebDriverException;
import org.luxor.commons.selenium.listener.SeleniumEventListener;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import javax.annotation.PreDestroy;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.Collection;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * SeleniumClient 实现类
 *
 * @author Mr.yan  @date 2022/6/10
 */
public class SeleniumClientImpl implements SeleniumClient {

    private final SeleniumDriverCreator seleniumDriverCreator = new SeleniumDriverCreator();

    private final static ConcurrentHashMap<String, SeleniumDriver> DRIVER_STORAGE = new ConcurrentHashMap(8);

    private final static long DEFAULT_SLEEP_TIMEOUT = 500;

    private final ThreadPoolTaskExecutor asyncExecutor;

    private final SeleniumProperties properties;

    private final Collection<SeleniumEventListener> listeners;


    public SeleniumClientImpl(SeleniumProperties properties, Collection<SeleniumEventListener> listeners) {
        this.properties = properties;
        this.listeners = listeners;
        this.asyncExecutor = initThreadPool(properties);
    }

    private ThreadPoolTaskExecutor initThreadPool(SeleniumProperties properties) {
        ThreadPoolTaskExecutor threadPool = new ThreadPoolTaskExecutor();
        // 设置核心线程数
        threadPool.setCorePoolSize(properties.getSessionMin());
        // 设置最大线程数
        threadPool.setMaxPoolSize(properties.getSessionMax());
        // 设置队列容量
        threadPool.setQueueCapacity(20);
        // 设置线程空闲时间（秒）
        threadPool.setKeepAliveSeconds(30);
        // 设置默认线程名称
        threadPool.setThreadNamePrefix("selenium-async-");
        // 设置拒绝策略
        threadPool.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        // 等待所有任务结束后再关闭线程池
        threadPool.setWaitForTasksToCompleteOnShutdown(true);
        // 初始化
        threadPool.initialize();
        return threadPool;
    }

    private void dispatcherInitEvent(SeleniumCommand command, Context context) {
        for (SeleniumEventListener listener : this.listeners) {
            listener.init(command, context);
        }
    }

    private void dispatcherBeforeEvent(SeleniumDriver webDriver, SeleniumCommand command, Context context) {
        for (SeleniumEventListener listener : this.listeners) {
            try {
                listener.beforeExecute(webDriver, command, context);
            } catch (Throwable throwable) {
                listener.onException(throwable, webDriver);
            }
        }
    }

    private <T> T run(SeleniumDriver webDriver, SeleniumCommand<T> command, Context context) {
        try {
            T result = command.execute(webDriver, context);
            return result;
        } catch (Throwable e) {
            if (ExecState.RUNNING.equals(context.getState()) && context.getExecuteCount() <= getProperties().getRetryNumMax()) {
                // 执行次数+1
                context.incrementAndExecuteCount();
                // 执行重试
                return run(webDriver, command, context);
            }
            throw e;
        }
    }

    private void dispatcherAfterEvent(SeleniumDriver webDriver, SeleniumCommand command, Context context) {
        for (SeleniumEventListener listener : this.listeners) {
            try {
                listener.afterExecute(webDriver, command, context);
            } catch (Throwable throwable) {
                listener.onException(throwable, webDriver);
            }
        }
    }

    /**
     * 执行浏览器驱动指令
     *
     * @param command 驱动指令
     * @return T
     */
    @Override
    public <T> T execute(SeleniumCommand<T> command) {
        SeleniumDriver webDriver = null;
        Context context = new Context();
        try {
            dispatcherInitEvent(command, context);
            context.execInitialize(getProperties(), command);
            webDriver = getDriver();
            context.execStart(webDriver, command);
            dispatcherBeforeEvent(webDriver, command, context);
            T result = run(webDriver, command, context);
            context.execCompleted();
            return result;
        } catch (Throwable e) {
            context.execFailed(e);
            throw e;
        } finally {
            dispatcherAfterEvent(webDriver, command, context);
            removeDriver(webDriver);
        }
    }


    /**
     * 异步执行浏览器驱动指令
     *
     * @param command 驱动指令
     * @return T
     */
    @Override
    public <T> Future<T> executeAsync(SeleniumCommand<T> command) {
        return asyncExecutor.submit(new SeleniumClientImpl.AsyncExecuteCallable<>(command));
    }

    class AsyncExecuteCallable<T> implements Callable<T> {
        private final SeleniumCommand<T> command;

        public AsyncExecuteCallable(SeleniumCommand<T> command) {
            this.command = command;
        }

        @Override
        public T call() {
            SeleniumDriver webDriver = null;
            Context context = new Context();
            try {
                dispatcherInitEvent(command, context);
                context.execInitialize(getProperties(), command);
                webDriver = getDriver();
                context.execStart(webDriver, command);
                dispatcherBeforeEvent(webDriver, this.command, context);
                T result = run(webDriver, this.command, context);
                context.execCompleted();
                return result;
            } catch (Throwable e) {
                context.execFailed(e);
                throw e;
            } finally {
                dispatcherAfterEvent(webDriver, this.command, context);
                removeDriver(webDriver);
            }
        }
    }

    public SeleniumProperties getProperties() {
        return this.properties;
    }

    /**
     * 获取浏览器驱动(SeleniumDriver)
     *
     * @return org.luxor.commons.selenium.component.SeleniumDriver
     */
    private SeleniumDriver getDriver() {
        Clock clock = Clock.systemDefaultZone();
        Duration timeout = Duration.ofSeconds(properties.getSessionCreateTimeoutSeconds());
        Instant end = clock.instant().plus(timeout);
        while (true) {
            // 限制webDriver最大可创建的数量，避免单台服务器运行过载
            if (DRIVER_STORAGE.size() < properties.getSessionMax()) {
                SeleniumDriver seleniumDriver = seleniumDriverCreator.creatorSeleniumDriver(properties, listeners);
                DRIVER_STORAGE.put(seleniumDriver.getSessionId(), seleniumDriver);
                return seleniumDriver;
            }

            // Check the timeout after evaluating the function to ensure conditions with a zero timeout can succeed.
            if (end.isBefore(clock.instant())) {
                throw new CreateWebDriverException("Unable to create a new WebDriver. " +
                        "The maximum number of webDrivers exceeded [" + properties.getSessionMax() + "].");
            }

            ThreadUtil.sleep(DEFAULT_SLEEP_TIMEOUT);
        }
    }

    /**
     * 移除浏览器驱动(SeleniumDriver)
     *
     * @param webDriver 浏览器驱动器
     */
    private void removeDriver(SeleniumDriver webDriver) {
        try {
            if (webDriver != null) {
                DRIVER_STORAGE.remove(webDriver.getSessionId());
                webDriver.quit();
            }
        } catch (Throwable e) {
            // .ignore
        }
    }

    @PreDestroy
    public void destroy() {
        DRIVER_STORAGE.values().forEach(SeleniumDriver::quit);
        DRIVER_STORAGE.clear();
    }
}
